#include "mt3d_cuda.h"

void mt3d::assemble_linear_system()
{
    std::clog << "Assembling kernel matrix ... ";
    size_t e, i, j;

    g3d_tet_.initiate(2);

    std::vector<std::vector<triplt>> omp_entries;
    omp_entries.resize(36);
    for (i = 0; i < 36; i++)
    {
        omp_entries[i].resize(elem_num_);
    }

    edge_setup e_set;
    double edge_cndt, edge_cndt_ref, aij_p1, aij_p2;
    cuda_cd tmp_entry;
    gctl::vertex3dc* tmp_vert[4];
    // Loop all elements
#pragma omp parallel for private (e, i, j, e_set, edge_cndt, aij_p1, aij_p2, tmp_entry, tmp_vert) schedule(guided)
    for (e = 0; e < elem_num_; e++)
    {
        tmp_vert[0] = mesh_elems_[e].vert[0];
        tmp_vert[1] = mesh_elems_[e].vert[1];
        tmp_vert[2] = mesh_elems_[e].vert[2];
        tmp_vert[3] = mesh_elems_[e].vert[3];

        get_element_conductivity(elem_markers_[e], edge_cndt);

        for (i = 0; i < 6; i++)
        {
            e_set.uv_id[0] = i;

            if (elem_edge_reversed_[e][i]) e_set.uv_type[0] = efem_sf::Reverse;
            else e_set.uv_type[0] = efem_sf::Normal;

            for (j = 0; j < 6; j++)
            {
                e_set.uv_id[1] = j;

                if (elem_edge_reversed_[e][j]) e_set.uv_type[1] = efem_sf::Reverse;
                else e_set.uv_type[1] = efem_sf::Normal;

                aij_p1 = g3d_tet_.integral(tmp_vert[0]->x, tmp_vert[1]->x, tmp_vert[2]->x, tmp_vert[3]->x, 
                    tmp_vert[0]->y, tmp_vert[1]->y, tmp_vert[2]->y, tmp_vert[3]->y, 
                    tmp_vert[0]->z, tmp_vert[1]->z, tmp_vert[2]->z, tmp_vert[3]->z, _Aij_part1, &e_set, this);

                aij_p2 = g3d_tet_.integral(tmp_vert[0]->x, tmp_vert[1]->x, tmp_vert[2]->x, tmp_vert[3]->x, 
                    tmp_vert[0]->y, tmp_vert[1]->y, tmp_vert[2]->y, tmp_vert[3]->y, 
                    tmp_vert[0]->z, tmp_vert[1]->z, tmp_vert[2]->z, tmp_vert[3]->z, _Aij_part2, &e_set, this);

                tmp_entry.x = aij_p1;
                tmp_entry.y = -1.0*omega_*_mu0*edge_cndt*aij_p2;

                omp_entries[i*6 + j][e].row = elem_edge_index_[e][i];
                omp_entries[i*6 + j][e].col = elem_edge_index_[e][j];
                omp_entries[i*6 + j][e].val = tmp_entry;
            }
        }
    }

    // Now we sort all triplts into a map
    size_t order;
    for (e = 0; e < elem_num_; e++)
    {
        for (i = 0; i < 36; i++)
        {
            order = edge_num_*(omp_entries[i][e].row) + omp_entries[i][e].col;

            if (triplt_map_.empty()) triplt_map_[order] = omp_entries[i][e];
            else
            {
                triplt_index_iter_ = triplt_map_.find(order);
                if (triplt_index_iter_ != triplt_map_.end()) // entriy exist
                {
                    tmp_entry = triplt_index_iter_->second.val;
                    tmp_entry = clcg_Zsum(tmp_entry, omp_entries[i][e].val);
                    triplt_index_iter_->second.val = tmp_entry;
                }
                else
                {
                    triplt_map_[order] = omp_entries[i][e];
                }
            }
        }
    }

/*
    triplt_index_iter_ = triplt_map_.begin();
    while (triplt_index_iter_ != triplt_map_.end())
    {
        std::cout << triplt_index_iter_->second.row << " " << triplt_index_iter_->second.col << " (" 
            << triplt_index_iter_->second.val.x << ", " << triplt_index_iter_->second.val.y << ")" << std::endl;
        triplt_index_iter_++;
    }
*/  

    nz_num_ = triplt_map_.size();
    h_rowIdxA_.resize(nz_num_);
    h_colIdxA_.resize(nz_num_);
    h_A_.resize(nz_num_);

    // Save the sparse matrix in the COO format
    int c = 0;
    for (triplt_index_iter_ = triplt_map_.begin(); triplt_index_iter_ != triplt_map_.end(); triplt_index_iter_++)
    {
        h_rowIdxA_[c] = triplt_index_iter_->second.row;
        h_colIdxA_[c] = triplt_index_iter_->second.col;
        h_A_[c] = triplt_index_iter_->second.val;
        c++;
    }

    for (i = 0; i < 36; i++)
    {
        omp_entries.clear();
    }
    triplt_map_.clear();

    // Now we initiate source terms
    h_B_TE_.resize(edge_num_, _zero);
    h_B_TM_.resize(edge_num_, _zero);

    // TE
    std::complex<double> tmp_inte;
    for (e = 0; e < elem_num_; ++e)
    {
        tmp_vert[0] = mesh_elems_[e].vert[0];
        tmp_vert[1] = mesh_elems_[e].vert[1];
        tmp_vert[2] = mesh_elems_[e].vert[2];
        tmp_vert[3] = mesh_elems_[e].vert[3];

        get_element_conductivity(elem_markers_[e], edge_cndt);
        get_element_conductivity(elem_markers_[e], edge_cndt_ref, true);
        edge_cndt -= edge_cndt_ref;

        if (edge_cndt != 0.0) // only calculate elements that have a residual conductivity
        {
            for (i = 0; i < 6; i++)
            {
                e_set.uv_id[0] = i;
                if (elem_edge_reversed_[e][i]) e_set.uv_type[0] = efem_sf::Reverse;
                else e_set.uv_type[0] = efem_sf::Normal;

                tmp_inte = omega_*_mu0*edge_cndt*
                    g3d_tet_.integral(tmp_vert[0]->x, tmp_vert[1]->x, tmp_vert[2]->x, tmp_vert[3]->x, 
                        tmp_vert[0]->y, tmp_vert[1]->y, tmp_vert[2]->y, tmp_vert[3]->y, 
                        tmp_vert[0]->z, tmp_vert[1]->z, tmp_vert[2]->z, tmp_vert[3]->z, _Bi, &e_set, this);

                h_B_TE_[elem_edge_index_[e][i]] = clcg_Zsum(h_B_TE_[elem_edge_index_[e][i]], cuCmul(lcg2cuda_complex(tmp_inte), _1i));
            }
        }
    }

    le_.set_to_TM_mode();
    le_.initialize();

    for (e = 0; e < elem_num_; ++e)
    {
        tmp_vert[0] = mesh_elems_[e].vert[0];
        tmp_vert[1] = mesh_elems_[e].vert[1];
        tmp_vert[2] = mesh_elems_[e].vert[2];
        tmp_vert[3] = mesh_elems_[e].vert[3];

        get_element_conductivity(elem_markers_[e], edge_cndt);
        get_element_conductivity(elem_markers_[e], edge_cndt_ref, true);
        edge_cndt -= edge_cndt_ref;

        if (edge_cndt != 0.0) // only calculate elements that have a residual conductivity
        {
            for (i = 0; i < 6; i++)
            {
                e_set.uv_id[0] = i;
                if (elem_edge_reversed_[e][i]) e_set.uv_type[0] = efem_sf::Reverse;
                else e_set.uv_type[0] = efem_sf::Normal;

                tmp_inte = omega_*_mu0*edge_cndt*
                    g3d_tet_.integral(tmp_vert[0]->x, tmp_vert[1]->x, tmp_vert[2]->x, tmp_vert[3]->x, 
                        tmp_vert[0]->y, tmp_vert[1]->y, tmp_vert[2]->y, tmp_vert[3]->y, 
                        tmp_vert[0]->z, tmp_vert[1]->z, tmp_vert[2]->z, tmp_vert[3]->z, _Bi, &e_set, this);

                h_B_TM_[elem_edge_index_[e][i]] = clcg_Zsum(h_B_TM_[elem_edge_index_[e][i]], cuCmul(lcg2cuda_complex(tmp_inte), _1i));
            }
        }
    }

    std::clog << "done\n";

    return;
}
